import requests
import re
from selenium import webdriver
import time
from contextlib import closing


def get_home_ulrs(url):
    browser = webdriver.Chrome();
    browser.get(url)
    num = browser.find_element_by_xpath(
        '//*[@id="root"]/div/div[2]/div/div[4]/div[1]/div[1]/div[1]/span').text
    num = int(int(num)/15)+1
    urls = []
    for i in range(num):
        newList = re.findall('https\:\/\/www\.douyin\.com\/video\/[\d]*\?previous_page\=', browser.page_source)
        urls = [*urls,*newList]
        browser.execute_script('window.scrollTo(0, document.body.scrollHeight)')
        time.sleep(2)
    print(urls)
    print(set(urls))
    mainList =set(urls)
    browser.close()
    return mainList

def get_video(url):
    browser = webdriver.Chrome();
    browser.get(url)
    html = browser.find_element_by_xpath(
        '//*[@id="root"]/div/div[2]/div[1]/div[1]/div[1]/div[1]/div/div[1]/div[2]/video/source[1]')
    video_url = html.get_attribute('src')
    browser.close()
    return video_url


def download(url, name):
    with closing(requests.get(url=url, verify=False, stream=True)) as res:
        with open('video/{}.mp4'.format(name), 'wb') as fd:
            for chunk in res.iter_content(chunk_size=1024):
                if chunk:
                    fd.write(chunk)

if __name__ == '__main__':
    url = input('请输入个人主页地址')
    print('正在获取视频地址')
    urls = get_home_ulrs(url)
    print(urls)
    print('成功获取{}个,开始下载'.format(len(urls)))
    index = 1
    for url in urls:
        video_url = get_video(url)
        print('正在下载{}/{}'.format(len(urls),index))
        download(video_url, index)
        print('下载完成')
        index += 1
